#!/usr/bin/env python
# encoding: utf-8
# pylint: disable-all


'''
www.xbiquge.cc 网站爬虫
抓取，生成电子书原材料
'''

import os
import time
import json
import uuid
import requests
from requests.adapters import HTTPAdapter
# from requests.packages.urllib3.util.retry import Retry
from urllib3.util.retry import Retry
from ebooklib import epub


def ebook_generator(obj):
  assert 'name' in obj and obj['name']
  assert 'chapters' in obj
  assert len(obj['chapters']) > 0
  book_name = obj['name']
  chapters = obj['chapters']
  # generate epub book.
  book = epub.EpubBook()
  # set metadata
  book.set_identifier(uuid.uuid1().hex)
  book.set_title(book_name)
  book.set_language('en')
  # set author
  if 'author' in obj:
    book.add_author(obj['author'])
  else:
    book.add_author('BookSpider by suncc')
  # create chapters
  book.toc = []
  book.spine = ['nav']
  contents = []
  for _, chapter in enumerate(chapters):
    fname = chapter['fname']
    with open(fname, 'rt', encoding="utf-8") as fin:
      contents.append(fin.read())
  for pos, chapter in enumerate(chapters):
    title = chapter['title']
    fname = chapter['fname']
    file_name = 'article-%d.xhtml' % (pos + 1)
    c1 = epub.EpubHtml(title=title, file_name=file_name, lang='en')
    c1.content=contents[pos]
    link1 = epub.Link(file_name, title, 'INDEX')
    book.add_item(c1)
    book.spine.append(c1)
    book.toc.append(link1)
  # add navigation files
  book.add_item(epub.EpubNcx())
  book.add_item(epub.EpubNav())
  # create epub file
  out_name = 'book/%s.epub' % book_name
  epub.write_epub(out_name, book, {})


class BookSpider:
  def __init__(self):
    self.session = requests.Session()
    self.prepared = False

  def user_agent(self):
    return 'Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:72.0) Gecko/20100101 Firefox/72.0'

  def referer(self):
    return ''

  def page_session(self, retries=3, backoff_factor=0.3, status_forcelist=(500, 502, 504)):
    session = self.session
    retry = Retry(
      total=retries,
      read=retries,
      connect=retries,
      backoff_factor=backoff_factor,
      status_forcelist=status_forcelist,
    )
    adapter = HTTPAdapter(max_retries=retry)
    session.mount('http://', adapter)
    session.mount('https://', adapter)
    return self.session

  def book_name(self):
    return 'MyBook'

  def need_sleep(self):
    return True

  def sections(self):
    '''
    返回所有需要的链接列表
    '''
    return []

  def download_url(self, url):
    headers = {
      'User-Agent': self.user_agent(),
    }
    referer = self.referer()
    if referer:
      headers['Referer'] = referer
    return self.page_session().get(url, headers=headers, timeout=15)

  def get_page(self, pos, url):
    '''
    下载某个链接，后期可能需要代理、伪装、识别等高端手段
    '''
    rsp = self.download_url(url)
    if rsp.status_code != 200:
      print('%d download %s failed.' % (pos, url))
      return None
    if self.need_sleep():
      time.sleep(2)
    return rsp.content

  def prepare(self):
    '''
    准备工作，子类重写
    '''

  def process_page(self, pos, url, page):
    '''
    处理原始 html 文档，返回 dict 对象，至少包含
    title, body 两部分
    '''
    return {
      "title": '',
      "body": '',
    }

  def process_section(self, pos, url, sections):
    return {
      "title": '',
      "body": '',
    }

  def save_page(self, pos, page):
    '''
    将书籍数据保存到本地文件系统
    '''
    body = page.get('body')
    bName = self.book_name()
    fname = 'book/%s/%03d.html' % (bName, pos)
    with open(fname, 'wb') as fout:
      fout.write(body)
    return fname

  def _load_book_data(self):
    default_obj = {
      "name": self.book_name(),
      "chapters": [],
    }
    file_name = 'book/%s.json' % self.book_name()
    if not os.path.exists(file_name):
      return default_obj
    with open(file_name, 'r+') as fIn:
      content = fIn.read()
    if not content:
      return default_obj
    else:
      return json.loads(content)

  def _dump_book_data(self, obj):
    file_name = 'book/%s.json' % self.book_name()
    with open(file_name, 'w+') as fout:
      fout.write(json.dumps(obj))

  def _update_book_data(self, pos, title, fname):
    obj = self._load_book_data()
    obj['chapters'].append({
      'pos': pos,
      'title': title,
      'fname': fname
    })
    self._dump_book_data(obj)

  def fetch_section(self, pos, url):
    '''
    若每一章节分成多个不同页面，子类复写重新处理
    '''
    return self.get_page(pos, url)

  def section_handle(self, pos, url):
    total = self.sections()
    print('准备开始抓取 %s 共 %d 章' % (self.book_name(), len(total)))
    try:
      page = self.fetch_section(pos, url)
    except requests.Timeout:
      print('%d download timeout' % pos)
      return
    if not page:
      return
    # fetch section 可能返回一系列页面
    if isinstance(page, list):
      page_obj = self.process_section(pos, url, page)
    elif isinstance(page, (bytes, str)):
      page_obj = self.process_page(pos, url, page)
    if not page_obj:
      return
    # 已处理为单个章节对象了
    file_name = self.save_page(pos, page_obj)
    title = page_obj['title']
    self._update_book_data(pos, title, file_name)
    print('%03d %s completed.' % (pos, title))

  def fix_item(self, pos):
    '''
    如果电子书某一章下载失败，可使用此方法单独补全, pos 从1起
    '''
    self.execute_prepare()
    sections = self.sections()
    print('修正 %s 第 %d 章' % (self.book_name(), pos))
    url = sections[pos - 1]
    self.section_handle(pos, url)

  def crawl_from(self, pos):
    '''
    断点续传
    章节从1开始算起
    '''
    self.execute_prepare()
    try:
      return self._run_from(pos)
    except KeyboardInterrupt:
      print('已终止！')

  def crawl(self):
    self.execute_prepare()
    try:
      return self._run()
    except KeyboardInterrupt:
      print('已终止！')

  def make_epub(self):
    '''
    直接生成 epub 格式的电子书
    抓取过程与电子书生成过程应隔绝开
    抓取后，应生成数据文件，为简化，可使用json
    使用数据文件生成 epub
    '''
    self.execute_prepare()
    obj = self._load_book_data()
    ebook_generator(obj)

  def execute_prepare(self):
    if not self.prepared:
      self.prepare()
      self.prepared = True

  def _create_book(self):
    if not os.path.exists('book'):
        os.mkdir('book')

  def _run_from(self, pos):
    sections = self.sections()
    tasks = sections[pos-1:]
    for idx, url in enumerate(tasks):
      newPos = idx + 1 + (pos - 1)
      self.section_handle(newPos, url)

  def _run(self):
    bName = self.book_name()
    full_path = 'book/%s' % bName
    if os.path.exists(full_path):
      print('书籍目录已存在')
      return None
    try:
      self._create_book()
      os.mkdir('book/%s' % bName)
    except OSError as e:
      print('书籍目录创建失败')
      print(e)
      return None
    self._run_from(1)
